import os
import asyncio
import zipfile
from pathlib import Path
from typing import List, Dict, Any
import json
import uuid

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from docx import Document
import openai
from dotenv import load_dotenv

load_dotenv()

app = FastAPI(title="大模型矩阵·智能插图生成器", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件服务
app.mount("/static", StaticFiles(directory="static"), name="static")

# 配置目录
UPLOAD_DIR = Path("uploads")
RESULTS_DIR = Path("results")
UPLOAD_DIR.mkdir(exist_ok=True)
RESULTS_DIR.mkdir(exist_ok=True)

# 全局任务状态存储
processing_tasks: Dict[str, Dict[str, Any]] = {}

class DocumentProcessor:
    def __init__(self):
        openai.api_key = os.getenv("OPENAI_API_KEY")
        if not openai.api_key:
            raise ValueError("请设置环境变量 OPENAI_API_KEY")

    async def extract_paragraphs(self, file_path: Path) -> List[Dict[str, str]]:
        """提取文档段落"""
        doc = Document(file_path)
        paragraphs = []
        
        for i, para in enumerate(doc.paragraphs):
            if para.text.strip():
                paragraphs.append({
                    "id": str(i),
                    "text": para.text.strip(),
                    "type": "paragraph"
                })
        
        return paragraphs

    async def generate_illustration(self, paragraph: Dict[str, str], task_id: str) -> Dict[str, str]:
        """为段落生成插图"""
        try:
            # 使用GPT生成图像描述
            response = openai.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "你是一个专业的插图设计师，请根据文本内容生成简洁的图像描述，适合用于AI图像生成。"},
                    {"role": "user", "content": f"请为以下文本生成一个简洁的插图描述（50字以内）：{paragraph['text']}"}
                ],
                max_tokens=100
            )
            
            image_description = response.choices[0].message.content
            
            # 使用DALL-E生成图像
            image_response = openai.images.generate(
                model="dall-e-3",
                prompt=image_description,
                size="1024x1024",
                quality="standard",
                n=1
            )
            
            return {
                "paragraph_id": paragraph["id"],
                "image_url": image_response.data[0].url,
                "description": image_description,
                "status": "completed"
            }
            
        except Exception as e:
            return {
                "paragraph_id": paragraph["id"],
                "error": str(e),
                "status": "failed"
            }

    async def process_document(self, file_path: Path, task_id: str) -> Dict[str, Any]:
        """处理单个文档"""
        try:
            # 提取段落
            paragraphs = await self.extract_paragraphs(file_path)
            
            # 更新任务状态
            processing_tasks[task_id]["total_paragraphs"] = len(paragraphs)
            processing_tasks[task_id]["processed_paragraphs"] = 0
            
            # 并行生成插图
            results = []
            for para in paragraphs:
                if processing_tasks[task_id].get("cancelled", False):
                    break
                    
                result = await self.generate_illustration(para, task_id)
                results.append(result)
                
                # 更新进度
                processing_tasks[task_id]["processed_paragraphs"] += 1
                processing_tasks[task_id]["progress"] = (
                    processing_tasks[task_id]["processed_paragraphs"] / 
                    processing_tasks[task_id]["total_paragraphs"] * 100
                )
            
            # 创建结果包
            result_dir = RESULTS_DIR / task_id
            result_dir.mkdir(exist_ok=True)
            
            # 保存结果
            result_data = {
                "original_filename": file_path.name,
                "paragraphs": paragraphs,
                "illustrations": results,
                "status": "completed"
            }
            
            with open(result_dir / "result.json", "w", encoding="utf-8") as f:
                json.dump(result_data, f, ensure_ascii=False, indent=2)
            
            processing_tasks[task_id]["status"] = "completed"
            processing_tasks[task_id]["result_path"] = str(result_dir)
            
            return result_data
            
        except Exception as e:
            processing_tasks[task_id]["status"] = "failed"
            processing_tasks[task_id]["error"] = str(e)
            raise e

processor = DocumentProcessor()

@app.post("/upload")
async def upload_files(files: List[UploadFile] = File(...)):
    """上传多个Word文档"""
    uploaded_files = []
    
    for file in files:
        if not file.filename.endswith('.docx'):
            raise HTTPException(status_code=400, detail=f"文件 {file.filename} 不是Word文档")
        
        task_id = str(uuid.uuid4())
        file_path = UPLOAD_DIR / f"{task_id}_{file.filename}"
        
        with open(file_path, "wb") as f:
            content = await file.read()
            f.write(content)
        
        # 初始化任务状态
        processing_tasks[task_id] = {
            "filename": file.filename,
            "status": "pending",
            "progress": 0,
            "total_paragraphs": 0,
            "processed_paragraphs": 0,
            "file_path": str(file_path)
        }
        
        uploaded_files.append({
            "task_id": task_id,
            "filename": file.filename
        })
    
    return {"uploaded_files": uploaded_files}

@app.post("/process/{task_id}")
async def process_document(task_id: str, background_tasks: BackgroundTasks):
    """开始处理文档"""
    if task_id not in processing_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if processing_tasks[task_id]["status"] != "pending":
        raise HTTPException(status_code=400, detail="任务状态无效")
    
    processing_tasks[task_id]["status"] = "processing"
    
    file_path = Path(processing_tasks[task_id]["file_path"])
    
    # 后台处理
    background_tasks.add_task(processor.process_document, file_path, task_id)
    
    return {"message": "处理已开始", "task_id": task_id}

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    """获取处理状态"""
    if task_id not in processing_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    return processing_tasks[task_id]

@app.get("/download/{task_id}")
async def download_results(task_id: str):
    """下载处理结果"""
    if task_id not in processing_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if processing_tasks[task_id]["status"] != "completed":
        raise HTTPException(status_code=400, detail="处理未完成")
    
    result_dir = Path(processing_tasks[task_id]["result_path"])
    zip_path = RESULTS_DIR / f"{task_id}_results.zip"
    
    # 创建压缩包
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for file in result_dir.rglob("*"):
            if file.is_file():
                zipf.write(file, file.relative_to(result_dir))
    
    return FileResponse(
        path=zip_path,
        filename=f"{task_id}_results.zip",
        media_type='application/zip'
    )

@app.delete("/cancel/{task_id}")
async def cancel_task(task_id: str):
    """取消处理任务"""
    if task_id not in processing_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    processing_tasks[task_id]["cancelled"] = True
    processing_tasks[task_id]["status"] = "cancelled"
    
    return {"message": "任务已取消"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
